/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef UTILS_TRIE_H
#define UTILS_TRIE_H

#include <cassert>
#include <vector>

#include "utils/log.h"

namespace profiler {

// A simple, non thread-safe Trie (prefix tree) backed by a vector.
// Note that the first node (index = 0) exists by default and contains
// A default value for the templated type. Currently only the insert operation
// is supported.
template <class T>
class Trie {
 public:
  Trie() : total_insert_count_(0), unique_path_count_(0) {
    // Default root node. TODO - currently assume T default constructible.
    nodes_.emplace_back(T(), -1);
  }

  // Attempt to insert a path into the prefix tree. On finish, returns a pair
  // containing the index of the leaf node and a bool value - True indicates
  // the node as generated by |values| is either new or marked as a leaf node
  // for the first time. False indicates the same path existed and the leaf
  // node already marked previously.
  std::pair<int, bool> insert(const std::vector<T>& values) {
    assert(values.size() > 0);
    bool new_leaf = false;

    // Start from root
    int index = 0;

    for (const auto& value : values) {
      bool found = false;
      for (auto child_index : nodes_.at(index).children_indices_) {
        if (value == nodes_.at(child_index).value_) {
          found = true;
          index = child_index;
          break;
        }
      }

      // Insert new node if necessary
      if (!found) {
        int new_index = nodes_.size();
        nodes_.emplace_back(value, index);
        nodes_.at(index).children_indices_.push_back(new_index);
        index = new_index;
      }
    }

    total_insert_count_++;
    // Mark the node at index as a leaf if it isn't already.
    if (!nodes_.at(index).is_leaf_) {
      nodes_.at(index).is_leaf_ = true;
      unique_path_count_++;
      new_leaf = true;
    }

    return std::make_pair(index, new_leaf);
  }

  // Returns the parent's position of the node specified by index.
  int getParent(int index) const { return nodes_.at(index).parent_index_; }

  // Returns the values of the node specified by index.
  const T& getValue(int index) const { return nodes_.at(index).value_; }

  // Returns whether the node specified by index is a leaf.
  bool isLeaf(int index) const { return nodes_.at(index).is_leaf_; }

  // Number of nodes in the Trie - including the root node.
  int length() const { return nodes_.size(); }

  // Bookkeeping - print tracking stats
  void PrintStats() const {
    Log::V(">> TotalNodes:%d, TotalInserts:%d, UniquePaths:%d", length(),
           total_insert_count_, unique_path_count_);
  }

 private:
  // An internal structure representing a node in the Trie.
  template <class U>
  struct TrieNode {
    TrieNode(const U& value, int parent_index)
        : value_(value), parent_index_(parent_index), is_leaf_(false) {}

    std::vector<int> children_indices_;
    const U value_;
    const int parent_index_;
    bool is_leaf_;
  };

  std::vector<TrieNode<T>> nodes_;
  int total_insert_count_;
  int unique_path_count_;
};

}  // namespace profiler

#endif  // UTILS_TRIE_H
